﻿#include "DoseStat.h"

#include <algorithm>
#include <cmath>
#include <cstdlib>
#include <functional>
#include <iostream>
#include <numeric>
#include <stdexcept>
#include <string>

/// 控制是否打印
#if 0
#    define MYDBG printf
#else
#    define MYDBG(...)
#endif

namespace dosestat
{

    // 计算DVH, dose-volume histogram (X-Y)
    DoseStatRet Accumulate(const std::vector<double>& voxelVolume, const std::vector<double>& voxelDose, bool doExactAccumulation)
    {
        MYDBG("doExactAccumulation: %d\n", doExactAccumulation);
        DoseStatRet result;

        // 预分配可能的内存以减少动态扩容; dose->volume
        const size_t numVoxels = voxelVolume.size();
        std::vector<std::pair<double, double>> doseVolumePairs;
        doseVolumePairs.reserve(numVoxels);

        /// 过滤体积为0的体素，同时直接创建 剂量->体积 对
        for (size_t i = 0; i < numVoxels; ++i)
        {
            if (voxelVolume[i] > 0)
            {
                doseVolumePairs.emplace_back(std::max(0.0, voxelDose[i]), voxelVolume[i]);
            }
        }
        doseVolumePairs.shrink_to_fit();
        // number of positive voxels
        const size_t num_pstv = doseVolumePairs.size();

        ///==== 空数据的快速处理
        if (num_pstv == 0)
        {
            /// ticks 最少为 2 => 一个 bin
            /// volume; 首哨兵为总体积, 末哨兵为 0.0
            /// dose; 首哨兵为 min dose, 末哨兵为 max dose
            result.accumulatedRelativeVolumes = {1.0, 0.0};
            result.binnedDoseValues = {0.0, 0.0};

            return result;
        }

        /// 根据数据量和doExactAccumulation参数选择合适的算法
        constexpr size_t EXACT_THRESHOLD = 5000;
        constexpr int MAX_BINS = 1001;
        constexpr int MIN_BINS = 2;
        constexpr double doseMinWidth = 1e-6;  // dose bin 最小分辨率
        const bool useExactMethod = doExactAccumulation || (num_pstv < EXACT_THRESHOLD);

        do
        {
            ///============================= 精确计算：
            /// 直接排序 {剂量->体积}, 先按照剂量排序，然后按照体积排序
            if (useExactMethod)
            {
                // 避免额外的索引数组和多次拷贝
                std::sort(doseVolumePairs.begin(), doseVolumePairs.end());

                // 准备结果数组
                result.binnedDoseValues.resize(num_pstv + 2, 0.0);
                result.accumulatedRelativeVolumes.resize(num_pstv + 2, 0.0);

                // 填充排序后的结果
                for (size_t i = 0; i < num_pstv; ++i)
                {
                    result.binnedDoseValues[i + 1] = doseVolumePairs[i].first;
                    result.accumulatedRelativeVolumes[i + 1] = doseVolumePairs[i].second;
                }

                break;  // skip all the rest
            }

            ///============================= 优化的粗略计算：
            /// 基于统计特征选择更合理的分箱数量; 提取剂量数据进行统计分析
            std::vector<double> doses;
            doses.reserve(num_pstv);
            double totalVolume = 0.0;

            for (const auto& [dose, volume] : doseVolumePairs)
            {
                doses.push_back(dose);
                totalVolume += volume;
            }

            /// 计算剂量范围
            auto [minDoseIt, maxDoseIt] = std::minmax_element(doses.begin(), doses.end());
            double minDose = std::max(0.0, *minDoseIt);
            double maxDose = *maxDoseIt;
            MYDBG("[my]minDose: %f, maxDose: %f, num_pstv: %zu\n", minDose, maxDose, num_pstv);

            /// 动态计算合适的分箱数
            const double doseRange = maxDose - minDose;

            ///============================= 粗略; 当剂量范围<=最小宽度, 用最小分箱数
            if (doseRange <= doseMinWidth)
            {
                char errMsg[256];
                sprintf(errMsg,
                        "Error: dose range to small, range -> %f, min -> %f, dose -> %f; while voxels number too much -> (%zu); check input data.\n",
                        doseRange, minDose, maxDose, num_pstv);
                throw std::invalid_argument(errMsg);
#if 0
            result.accumulatedRelativeVolumes = {1.0, 0.0};
            result.binnedDoseValues = {minDose, maxDose};
            return result;
#endif
            }

            ///============================= 当剂量范围>最小宽度，动态计算分箱数
            // 根据剂量范围和体素数量自适应确定分箱数
            int numBins = MIN_BINS;
            // 计算理论分箱数：每doseMinWidth剂量一个分箱; 加1.0 确保覆盖整个剂量范围
            const double theoreticalBins = doseRange / doseMinWidth + 1.0;
            // 将理论分箱数转换为整数，并限制在最小和最大分箱数之间
            const auto calculatedBins = static_cast<int>(theoreticalBins);
            numBins = std::max(MIN_BINS, std::min(MAX_BINS, calculatedBins));

            /// 准备分箱结果
            // 0 位置, n+1 位置, 两个哨兵
            constexpr size_t tickOfst = 1;
            result.binnedDoseValues.resize(numBins + 2, 0.0);
            result.accumulatedRelativeVolumes.resize(numBins + 2, 0.0);
            const double binWidth = doseRange / (numBins - 1);

            // 填充分箱剂量值
            for (int j = 0; j < numBins; ++j)
            {
                result.binnedDoseValues[j + tickOfst] = minDose + j * binWidth;
            }

            /// 统计每个分箱的体积
            for (const auto& [dose, volume] : doseVolumePairs)
            {
                if (dose <= minDose)  // 哨兵 0
                {
                    result.accumulatedRelativeVolumes[tickOfst] += volume;
                    continue;
                }
                if (dose >= maxDose)  // 哨兵 n+1
                {
                    //result.accumulatedRelativeVolumes[numBins + tickOfst] += volume;
                    continue;
                }
                // 落点取 bin 的右边界, 对应高的 dose
                const int binIndex = static_cast<int>((dose - minDose) / binWidth + 0.5);
                result.accumulatedRelativeVolumes[binIndex + tickOfst] += volume;
            }

            break;

        } while (0);

        // 计算累积相对体积
        const auto numBins = static_cast<int>(result.accumulatedRelativeVolumes.size()) - 2;
        double totalVol = 0.0;
        for (int i = numBins; i > 0; --i)
        {
            /// i-1 的累积体积 + 0.5 的当前体积, 偏低; counter 上面取右侧偏高;
            const double vol = result.accumulatedRelativeVolumes[i];
            result.accumulatedRelativeVolumes[i] = totalVol + 0.5 * vol;
            totalVol += vol;
        }

        /// 哨兵 规整
        result.accumulatedRelativeVolumes[0] = totalVol;
        result.accumulatedRelativeVolumes.back() = 0.0;
        result.binnedDoseValues[0] = result.binnedDoseValues[1];
        result.binnedDoseValues.back() = result.binnedDoseValues[numBins];

        /// 归一化
        if (totalVol > 0.0)
        {
            for (auto& volume : result.accumulatedRelativeVolumes)
            {
                volume /= totalVol;
            }
        }

        return result;
    }

    //xxxxxxxxxxxxxxxxxxxxxxxxxxxx 获取特定体积百分比对应的剂量
    double GetMetric(double doseAtVolume, const std::vector<double>& binnedDoseValues, const std::vector<double>& accumulatedRelativeVolumes)
    {
        /// 参数校验
        if (accumulatedRelativeVolumes.empty() || binnedDoseValues.empty())
        {
            MYDBG("Warning: Empty input vectors provided to GetMetric\n");
            return 0.0;
        }

        /// 边界检查优化
        if (doseAtVolume >= accumulatedRelativeVolumes[0])
        {
            return binnedDoseValues[0];
        }
        if (doseAtVolume <= accumulatedRelativeVolumes.back())
        {
            return binnedDoseValues.back();
        }

        /// 二分查找, 返回 所有等于 value 的元素范围
        // [](double l1, double l2) { return l1 > l2; }
        auto [first, last] =
            std::equal_range(accumulatedRelativeVolumes.begin(), accumulatedRelativeVolumes.end(), doseAtVolume, std::greater<double>{});

        ///================= 存在恰好等于 dose 的情况 [v1,v2,...], rarely
        if (first != last)
        {
            // 计算first到last范围内binnedDoseValues的平均值
            size_t count = 0;
            double sumDoses = 0.0;
            for (auto it = first; it != last; ++it)
            {
                const size_t index = it - accumulatedRelativeVolumes.begin();
                sumDoses += binnedDoseValues[index];
                count++;
            }
            const double avgDose = sumDoses / count;
            MYDBG("Volume -> %f, matching values -> (count: %zu), average dose -> %f \n", doseAtVolume, count, avgDose);
            return avgDose;
        }

        ///================= 不存在等于 dose 的情况, 插值; lower_bound == upper_bound
        const size_t ind = first - accumulatedRelativeVolumes.begin();
        MYDBG("Target dose -> %f, found index -> %zu\n", doseAtVolume, ind);
        ///------------- 边界检查
        if (ind == 0)
        {
            return binnedDoseValues[0];
        }
        if (ind >= accumulatedRelativeVolumes.size())
        {
            return binnedDoseValues.back();
        }
        size_t ind1 = ind - 1;  // little dose; big volume
        size_t ind2 = ind;      // big dose; small volume

        ///------------- 如果体积接近，返回平均剂量
        const double volDiff = accumulatedRelativeVolumes[ind2] - accumulatedRelativeVolumes[ind1];
        ///
        if (std::abs(volDiff) <= 1e-10)
        {
            MYDBG("Average; dose -> %f, index -> %zu\n", doseAtVolume, ind);
            return 0.5 * (binnedDoseValues[ind2] + binnedDoseValues[ind1]);
        }

        ///------------- 线性插值
        MYDBG("Linear interpolation; dose -> %f, index -> %zu\n", doseAtVolume, ind);
        double doseDiff = binnedDoseValues[ind2] - binnedDoseValues[ind1];
        double metric = binnedDoseValues[ind1] + (doseAtVolume - accumulatedRelativeVolumes[ind1]) * doseDiff / volDiff;
        return metric;
    }

    //xxxxxxxxxxxxxxxxxxxxxxxxxxx 计算ROI内的剂量统计指标
    double GetDoseStatRoiInterp(const std::string& dose_type,
                                const std::vector<double>& binnedDoseValues,
                                const std::vector<double>& accumulatedRelativeVolumes,
                                const std::vector<double>& voxelVolume,
                                const std::vector<double>& voxelDose,
                                double doseAtVolume)
    {
        ///================= 计算平均剂量时，可以提前检查数组大小是否匹配
        if (voxelVolume.size() != voxelDose.size())
        {
            throw std::runtime_error("voxelVolume and voxelDose arrays have different sizes");
            return 0.0;
        }

        ///================= 使用哈希值或枚举类型替代字符串比较，提高性能
        if (dose_type == "average")
        {
            double totalVolume = 0.0;
            double weightedDose = 0.0;
            for (size_t i = 0; i < voxelVolume.size(); ++i)
            {
                totalVolume += voxelVolume[i];
                weightedDose += voxelDose[i] * voxelVolume[i];
            }
            if (totalVolume <= 0.0)
            {
                throw std::runtime_error("Prescription ROI volume is zero. Cannot scale to average dose.");
            }

            return weightedDose / totalVolume;
        }

        ///================= 提取目标体积百分比
        double targetDoseAtVolume;
        // 使用更高效的查找方式，例如unordered_map或者switch-case(如果使用枚举)
        if (dose_type == "D_98")
        {
            targetDoseAtVolume = 0.98;
        }
        else if (dose_type == "D_95")
        {
            targetDoseAtVolume = 0.95;
        }
        else if (dose_type == "D_50")
        {
            targetDoseAtVolume = 0.5;
        }
        else if (dose_type == "D_2")
        {
            targetDoseAtVolume = 0.02;
        }
        else if (dose_type == "DoseAtVolume")
        {
            targetDoseAtVolume = doseAtVolume;
        }
        else
        {
            throw std::runtime_error("Prescription type not supported. Only 'average', 'D_98', 'D_95', 'D_50', 'D_2', 'DoseAtVolume' are supported");
        }

        // 调用优化后的get_metric函数
        return GetMetric(targetDoseAtVolume, binnedDoseValues, accumulatedRelativeVolumes);
    }
}  // namespace dosestat
